#include "tracknetv3_trt.h"
#include "cuda_utils.h"
#include <fstream>
#include <cassert>
#include <algorithm>

TrackNetV3_TRT::TrackNetV3_TRT() 
    : runtime_(nullptr), engine_(nullptr), context_(nullptr),
      input_buffer_(nullptr), output_buffer_(nullptr),
      input_index_(0), output_index_(0),
      max_batch_size_(1), input_width_(0), input_height_(0), input_channels_(0),
      output_width_(0), output_height_(0), output_channels_(0),
      input_size_(0), output_size_(0), stream_(nullptr) {
    
    std::cout << "Creating TrackNetV3_TRT instance..." << std::endl;
}

TrackNetV3_TRT::~TrackNetV3_TRT() {
    std::cout << "Destroying TrackNetV3_TRT instance..." << std::endl;
    destroy();
    if (stream_) {
        cudaStreamDestroy(stream_);
        std::cout << "CUDA stream destroyed" << std::endl;
    }
}

void TrackNetV3_TRT::printDims(const nvinfer1::Dims& dims, const std::string& name) {
    std::cout << name << " dimensions: [";
    for (int i = 0; i < dims.nbDims; ++i) {
        std::cout << dims.d[i];
        if (i < dims.nbDims - 1) std::cout << ", ";
    }
    std::cout << "]" << std::endl;
}

bool TrackNetV3_TRT::init(const std::string& engine_path) {
    std::cout << "Initializing TrackNetV3 with engine: " << engine_path << std::endl;
    
    // 检查文件是否存在
    std::ifstream engine_file(engine_path, std::ios::binary);
    if (!engine_file) {
        std::cerr << "Error: Unable to open engine file: " << engine_path << std::endl;
        return false;
    }
    
    engine_file.seekg(0, std::ifstream::end);
    size_t size = engine_file.tellg();
    engine_file.seekg(0, std::ifstream::beg);
    
    if (size == 0) {
        std::cerr << "Error: Engine file is empty: " << engine_path << std::endl;
        return false;
    }
    
    std::cout << "Engine file size: " << size << " bytes" << std::endl;
    
    std::unique_ptr<char[]> engine_data(new char[size]);
    engine_file.read(engine_data.get(), size);
    
    if (!engine_file) {
        std::cerr << "Error: Failed to read engine file" << std::endl;
        return false;
    }
    engine_file.close();
    
    // 创建运行时
    runtime_ = nvinfer1::createInferRuntime(logger_);
    if (!runtime_) {
        std::cerr << "Error: Failed to create TensorRT runtime" << std::endl;
        return false;
    }
    std::cout << "TensorRT runtime created successfully" << std::endl;
    
    // 反序列化引擎
    engine_ = runtime_->deserializeCudaEngine(engine_data.get(), size);
    if (!engine_) {
        std::cerr << "Error: Failed to deserialize CUDA engine" << std::endl;
        return false;
    }
    std::cout << "CUDA engine deserialized successfully" << std::endl;
    
    // 创建执行上下文
    context_ = engine_->createExecutionContext();
    if (!context_) {
        std::cerr << "Error: Failed to create execution context" << std::endl;
        return false;
    }
    std::cout << "Execution context created successfully" << std::endl;
    
    // 获取绑定信息
    int num_bindings = engine_->getNbBindings();
    std::cout << "Number of bindings: " << num_bindings << std::endl;
    
    for (int i = 0; i < num_bindings; ++i) {
        const char* name = engine_->getBindingName(i);
        bool is_input = engine_->bindingIsInput(i);
        nvinfer1::Dims dims = engine_->getBindingDimensions(i);
        
        std::cout << "Binding " << i << ": " << name 
                  << " (" << (is_input ? "input" : "output") << ")" << std::endl;
        printDims(dims, name);
        
        if (is_input) {
            input_index_ = i;
            // 解析输入维度
            if (dims.nbDims == 4) {
                max_batch_size_ = dims.d[0];
                input_channels_ = dims.d[1];
                input_height_ = dims.d[2];
                input_width_ = dims.d[3];
            } else {
                std::cerr << "Error: Unexpected input dimensions" << std::endl;
                return false;
            }
        } else {
            output_index_ = i;
            // 解析输出维度
            if (dims.nbDims == 4) {
                output_channels_ = dims.d[1];
                output_height_ = dims.d[2];
                output_width_ = dims.d[3];
            } else {
                std::cerr << "Error: Unexpected output dimensions" << std::endl;
                return false;
            }
        }
    }
    
    // 计算内存大小
    input_size_ = max_batch_size_ * input_channels_ * input_height_ * input_width_ * sizeof(float);
    output_size_ = max_batch_size_ * output_channels_ * output_height_ * output_width_ * sizeof(float);
    
    std::cout << "Input size: " << input_size_ << " bytes" << std::endl;
    std::cout << "Output size: " << output_size_ << " bytes" << std::endl;
    
    // 分配GPU内存
    CUDA_CHECK(cudaMalloc(&input_buffer_, input_size_));
    CUDA_CHECK(cudaMalloc(&output_buffer_, output_size_));
    
    // 创建CUDA流
    CUDA_CHECK(cudaStreamCreate(&stream_));
    
    std::cout << "GPU memory allocated successfully" << std::endl;
    std::cout << "TrackNetV3 initialization completed successfully!" << std::endl;
    std::cout << "Max batch size: " << max_batch_size_ << std::endl;
    std::cout << "Input: " << input_channels_ << "x" << input_height_ << "x" << input_width_ << std::endl;
    std::cout << "Output: " << output_channels_ << "x" << output_height_ << "x" << output_width_ << std::endl;
    
    return true;
}

bool TrackNetV3_TRT::preprocessFrames(const std::vector<cv::Mat>& frames, std::vector<float>& input_data) {
    if (frames.size() != 3) {
        std::cerr << "Error: TrackNetV3 requires exactly 3 input frames" << std::endl;
        return false;
    }
    
    int channel_size = input_height_ * input_width_;
    input_data.resize(input_channels_ * channel_size);
    
    for (int frame_idx = 0; frame_idx < 3; ++frame_idx) {
        const cv::Mat& frame = frames[frame_idx];
        
        // 调整尺寸
        cv::Mat resized;
        cv::resize(frame, resized, cv::Size(input_width_, input_height_));
        
        // 转换为浮点型并归一化
        cv::Mat normalized;
        resized.convertTo(normalized, CV_32FC3, 1.0 / 255.0);
        
        // 转换为RGB顺序
        cv::Mat rgb;
        cv::cvtColor(normalized, rgb, cv::COLOR_BGR2RGB);
        
        // 分割通道
        std::vector<cv::Mat> channels(3);
        cv::split(rgb, channels);
        
        // 将当前帧的3个通道放入对应的位置
        for (int c = 0; c < 3; ++c) {
            int channel_index = frame_idx * 3 + c;
            float* channel_data = input_data.data() + channel_index * channel_size;
            
            if (channels[c].total() != channel_size) {
                std::cerr << "Error: Channel size mismatch" << std::endl;
                return false;
            }
            
            memcpy(channel_data, channels[c].data, channel_size * sizeof(float));
        }
    }
    
    return true;
}

bool TrackNetV3_TRT::process(const std::vector<cv::Mat>& input_frames, cv::Mat& output_heatmap) {
    if (input_frames.size() < 3) {
        std::cerr << "Error: TrackNetV3 requires at least 3 consecutive frames" << std::endl;
        return false;
    }
    
    // 预处理输入帧
    std::vector<float> input_data;
    if (!preprocessFrames(input_frames, input_data)) {
        return false;
    }
    
    // 拷贝数据到GPU
    CUDA_CHECK(cudaMemcpyAsync(input_buffer_, input_data.data(), 
                              input_size_, cudaMemcpyHostToDevice, stream_));
    
    // 准备绑定数据
    void* bindings[2] = {input_buffer_, output_buffer_};
    
    // 执行推理
    bool success = context_->executeV2(bindings);
    if (!success) {
        std::cerr << "Error: Failed to execute inference" << std::endl;
        return false;
    }
    
    // 从GPU读取结果
    std::vector<float> output_data(output_channels_ * output_height_ * output_width_);
    CUDA_CHECK(cudaMemcpyAsync(output_data.data(), output_buffer_, 
                              output_size_, cudaMemcpyDeviceToHost, stream_));
    
    cudaStreamSynchronize(stream_);
    
    // 后处理
    postprocessOutput(output_data.data(), output_heatmap);
    
    return true;
}

void TrackNetV3_TRT::postprocessOutput(float* output, cv::Mat& heatmap) {
    heatmap.create(output_height_, output_width_, CV_32FC(output_channels_));
    
    if (output_channels_ == 1) {
        // 单通道热力图
        memcpy(heatmap.data, output, output_height_ * output_width_ * sizeof(float));
    } else {
        // 多通道热力图，取第一个通道
        std::vector<cv::Mat> channels(output_channels_);
        for (int c = 0; c < output_channels_; ++c) {
            float* channel_data = output + c * output_height_ * output_width_;
            channels[c] = cv::Mat(output_height_, output_width_, CV_32FC1, channel_data);
        }
        cv::merge(channels, heatmap);
    }
}

bool TrackNetV3_TRT::getBallPosition(const cv::Mat& heatmap, int& x, int& y, float threshold) {
    if (heatmap.empty()) {
        std::cerr << "Error: Empty heatmap" << std::endl;
        return false;
    }
    
    cv::Mat processed_heatmap;
    if (heatmap.channels() > 1) {
        // 多通道热力图，取第一个通道或合并
        std::vector<cv::Mat> channels;
        cv::split(heatmap, channels);
        processed_heatmap = channels[0];
    } else {
        processed_heatmap = heatmap;
    }
    
    // 找到热力图中的最大值位置
    cv::Point max_loc;
    double max_val;
    cv::minMaxLoc(processed_heatmap, nullptr, &max_val, nullptr, &max_loc);
    
    // 如果最大值低于阈值，则认为没有检测到球
    if (max_val < threshold) {
        x = 0;
        y = 0;
        return false;
    }
    
    x = max_loc.x;
    y = max_loc.y;
    return true;
}

bool TrackNetV3_TRT::destroy() {
    std::cout << "Cleaning up TrackNetV3 resources..." << std::endl;
    
    if (input_buffer_) {
        cudaFree(input_buffer_);
        input_buffer_ = nullptr;
    }
    
    if (output_buffer_) {
        cudaFree(output_buffer_);
        output_buffer_ = nullptr;
    }
    
    if (context_) {
        context_->destroy();
        context_ = nullptr;
    }
    
    if (engine_) {
        engine_->destroy();
        engine_ = nullptr;
    }
    
    if (runtime_) {
        runtime_->destroy();
        runtime_ = nullptr;
    }
    
    std::cout << "TrackNetV3 resources cleaned up successfully" << std::endl;
    return true;
}
